import os
import openslide
import tqdm
import time
from PIL import Image
import numpy as np
import cv2
import copy
import xml.etree.ElementTree as ET
import shutil
import argparse
# python3 automated_segmentaition_ndpi.py --ndpi_path '/media/alex/FA5EB5A15EB556DB1/uterus/2022.3.15子宫确诊/2023-03-15 20.50.55.ndpi'
parser = argparse.ArgumentParser(description='Ndpi to xml')
parser.add_argument('--ndpi_path', '-f', type=str, default=False, help='Load ndpi from a .ndpi file')
args = parser.parse_args()

ndpi_path = args.ndpi_path
png_path = './png_test' # png dir ,cropped by ndpi



if os.path.exists(png_path):
    shutil.rmtree(png_path )
    os.mkdir(png_path )
else:
    os.mkdir(png_path )
time1 = time.time()
slide = openslide.open_slide(ndpi_path)
print ('opentime :',time.time()-time1)
width, height = slide.dimensions
print (width,height)
index_x = 0
index_y = 0
index = 0

# crop ndpi ,get png images
while True:
    while True:
        tile_size_x = 1024 if index_x+1024<width else width-index_x+1
        tile_size_y = 1024 if index_y+1024<height else height-index_y+1

        time1 = time.time()
        region = slide.read_region((index_x, index_y), 0, (tile_size_x, tile_size_y))
        print ('read_region_time :',time.time()-time1)
        
        time1 = time.time()
        # simg1=os.path.join(png_path,'{}_0000.png'.format(str(index).zfill(4)))
        simg1=os.path.join(png_path,'{}_{}_{}_0000.png'.format(index_x,index_y,str(index).zfill(4)))
        region = region.convert('RGB')
        region.save(simg1)
        print ('save_time :',time.time()-time1)
        
        
        index +=1
        if index_x+1024<width:
            index_x +=924 
        else:
            index_x = 0
            break

    if index_y+1024<height:
            index_y +=924  
    else:
        index_y = 0
        break
slide.close()




from ultralytics import YOLO

# 加载一个模型
model = YOLO('/home/alex/Desktop/code/zjc_project/yolov8/runs/detect/train13/weights/best.pt')  # 从YAML建立并转移权重

results = model.predict('./png_test',imgsz=1024,conf=0.5)

# from ultralytics.engine.results

# load xml 
xml_path = './module.xml'

tree = ET.parse(xml_path)
root = tree.getroot()
annotation =  list(root.iter('Annotation'))[0]

# draw cells
regions  = list(annotation.iter('Regions'))[0]
region  = list(annotation.iter('Region'))[0]
Vertices = list(region.iter('Vertices'))[0]
vertex = list(Vertices.iter('Vertex'))[0]

root.remove(annotation)
annotation.remove(regions)
regions.remove(region)
region.remove(Vertices)
Vertices.remove(vertex)

ID = 1
for result in results:
    if len(result.boxes)>0:
        image_x = int(os.path.split(result.path)[-1].split('_')[0])
        image_y = int(os.path.split(result.path)[-1].split('_')[1])
        for xywh in result.boxes.xywh:
            x,y,w,h = int(xywh[0])+image_x,int(xywh[1])+image_y,int(xywh[2]),int(xywh[3])
            points = [
                [x,y],[x+w,y],[x+w,y+h],[x,y+h],[x,y-1]
            ]
            region_to_save = copy.deepcopy(region)
            Vertices_to_save = copy.deepcopy(Vertices )

            region_to_save.set('Id',str(ID))
            region_to_save.set('DisplayId',str(ID))

            for point in points:
                vertex_to_save = copy.deepcopy(vertex)
                vertex_to_save.set('X',str(point[0]))
                vertex_to_save.set('Y',str(point[1]))
                Vertices_to_save.append(vertex_to_save)
            region_to_save.append(Vertices_to_save)
            regions.append(region_to_save)
            ID += 1

annotation.append(regions)
root.append(annotation)

tree.write(ndpi_path[:-5]+'.xml')

